# coding:utf-8
from pathlib import Path
from PyPDF2 import PdfWriter, PdfReader
from reportlab.lib.units import cm
from reportlab.pdfgen import canvas
import reportlab.pdfbase.ttfonts
from customizeWindowPyfile.ProgressBarDialog import ProgressBar


def create_watermark(content, angle, alpha, color):
    # 制作水印文件
    file_name = 'function/Pdf/watermark.pdf'
    a = canvas.Canvas(file_name, pagesize=(50 * cm, 50 * cm))
    a.translate(0 * cm, 0 * cm)
    reportlab.pdfbase.pdfmetrics.registerFont(
        reportlab.pdfbase.ttfonts.TTFont('站酷高端黑', 'function/Pdf/font/圆体-简繁 细体.ttc'))
    a.setFont('站酷高端黑', 26)
    a.rotate(angle)  # 旋转角度
    a.setFillColorRGB(color[0], color[1], color[2])  # 选择颜色
    a.setFillAlpha(alpha)  # 透明度
    for i in range(0, 50, 5):
        for j in range(0, 50, 5):
            a.drawString(i * cm, j * cm, content)
    a.save()
    return file_name


def add_watermark(pdf_file_in, pdf_file_mark, pdf_file_out):
    # 添加水印文件到PDF文件中
    outputfile = PdfWriter()
    inputfile = PdfReader(pdf_file_in)
    pageCount = len(inputfile.pages)
    markfile = PdfReader(pdf_file_mark)
    # pdf文件与水印文件的合并
    for i in range(pageCount):
        # 在PDF文件主页加水印，形成新的PDF文件
        page = inputfile.pages[i]
        page.merge_page(markfile.pages[0])
        outputfile.add_page(page)
    with open(pdf_file_out, 'wb') as f_out:
        # 写入到输出文件中，形成新的文件
        outputfile.write(f_out)


def createWatermark4PDF(content, pdfPath, color=[139, 105, 105], angle=10, alpha=0.5):
    # 单个或多个PDF文件加水印处理
    color = [color[0] / 255, color[1] / 255, color[2] / 255]

    des_folder = None
    try:
        # 单个文件处理
        suffix = pdfPath.split('.')[1]
        file_list = [Path(pdfPath)]
        from_file = Path(pdfPath)
        if suffix == 'pdf':
            des_folder = Path(str(from_file.resolve()).strip('.pdf') + '水印')
    except IndexError:
        # 多个文件处理
        src_folder = Path(pdfPath)
        file_list = list(src_folder.glob('*.pdf*'))
        des_folder = src_folder / '水印'

    if des_folder is not None and not des_folder.exists():
        des_folder.mkdir(parents=True)

    totalTaskNum = len(file_list)
    doneTaskNum = 0
    currTaskNum = 1
    bar = ProgressBar()
    bar.setProcessOnTiTle(currTaskNum, totalTaskNum)
    bar.show()

    for pdf in file_list:
        # PDF添加水印
        pdf_file_in = str(pdf)
        pdf_file_mark = create_watermark(content, angle, alpha, color)
        pdf_file_out = str(des_folder / pdf.name)
        add_watermark(pdf_file_in, pdf_file_mark, pdf_file_out)

        # 转换进度条
        doneTaskNum = doneTaskNum + 1
        currTaskNum = currTaskNum + 1
        bar.setValue(100 * doneTaskNum / totalTaskNum, 100 * (doneTaskNum - 1) / totalTaskNum,
                     100 * totalTaskNum / totalTaskNum)


if __name__ == '__main__':
    color = [139, 105, 105]
    createWatermark4PDF('天才制造', 'D:\Desktop\功能测试\PDF文件')
